# 创建时间：2024/9/24下午11:09
# 实现AES的GCM模式加解密
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.backends import default_backend
import os
import abc


class Cipherer(metaclass=abc.ABCMeta):
    def __init__(self, key=None, nonce=None):
        self._key = None
        """秘钥"""
        self.key = key
        self._nonce = None
        """随机数"""
        self.nonce = nonce
        self._backend = default_backend()
        self._plaintext = None
        self._ciphertext = None
        self._tag = None

    @property
    def key(self):
        return self._key

    @key.setter
    def key(self, key):
        """设置秘钥, 如果值为空, 则生成一个128位随机数作为秘钥"""
        self._key = key if key else os.urandom(32)

    @property
    def nonce(self):
        return self._nonce

    @nonce.setter
    def nonce(self, nonce):
        """设置随机数, 如果值为空, 则生成一个96位随机数作为随机数"""
        self._nonce = nonce if nonce else os.urandom(12)

    @property
    def plaintext(self):
        """明文"""
        return self._plaintext

    @property
    def ciphertext(self):
        """密文"""
        return self._ciphertext

    @property
    def tag(self):
        return self._tag


class Encryptor(Cipherer):
    def __init__(self, key=None, nonce=None):
        super().__init__(key, nonce)
        self._encryptor = self._create_encryptor()

    @property
    def encryptor(self):
        return self._encryptor

    def _create_encryptor(self):
        """创建加密器"""
        cipher = Cipher(algorithms.AES(self.key), modes.GCM(self.nonce), backend=self._backend)
        encryptor = cipher.encryptor()
        return encryptor

    def rebuild_encryptor(self):
        """重建加密器"""
        cipher = Cipher(algorithms.AES(self.key), modes.GCM(self.nonce), backend=self._backend)
        encryptor = cipher.encryptor()
        self._encryptor = encryptor

    def encrypt(self, data):
        """
        对数据进行加密, 返回秘钥、随机数、tag、密文的四元组
        :param data: 要加密的数据
        :return: 返回秘钥、随机数、tag、密文的四元组
        """
        if not data:
            return None
        if not isinstance(data, bytes):
            raise TypeError(f"加密数据需要是bytes类型, 当前数据类型为{type(data)}")
        self._plaintext = data
        # 加密
        data = self.encryptor.update(data) + self.encryptor.finalize()
        # 获取加密后的认证标签
        tag = self.encryptor.tag
        self._ciphertext = data
        self._tag = tag
        return self.key, self.nonce, tag, data


class Decryptor(Cipherer):
    def __init__(self, key, nonce, tag, data=None):
        super().__init__(key, nonce)
        self._tag = tag
        self._ciphertext = data
        self._decryptor = self._create_decryptor()

    @Cipherer.key.setter
    def key(self, key):
        if not key:
            raise ValueError("解密器的秘钥不能为空")
        self._key = key

    @Cipherer.nonce.setter
    def nonce(self, nonce):
        if not nonce:
            raise ValueError("解密器的随机数不能为空")
        self._nonce = nonce

    @Cipherer.tag.setter
    def tag(self, tag):
        if not tag:
            raise ValueError("解密器的tag不能为空")
        self._tag = tag

    @property
    def decryptor(self):
        return self._decryptor

    def _create_decryptor(self):
        """创建解密器"""
        cipher = Cipher(algorithms.AES(self._key), modes.GCM(self._nonce, self._tag), backend=self._backend)
        decryptor = cipher.decryptor()
        return decryptor

    def rebuild_decryptor(self):
        """重建解密器"""
        cipher = Cipher(algorithms.AES(self.key), modes.GCM(self.nonce, self.tag), backend=self._backend)
        decryptor = cipher.decryptor()
        self._decryptor = decryptor

    def decrypt(self, data=None):
        """
        对数据进行解密, 解密后的数据
        :param data: 要解密的数据
        :return: 返回解密后的数据
        """
        if not data and not self._ciphertext:
            raise ValueError("未接收到数据")
        if not self._ciphertext:
            self._ciphertext = data
        if not data:
            data = self._ciphertext
        # 解密
        data = self.decryptor.update(data) + self.decryptor.finalize()
        # 获取加密后的认证标签
        self._plaintext = data
        return data


def test():
    encryptor = Encryptor()
    key, nonce, tag, data = encryptor.encrypt(b'abcdefghijklmnopqrstuvwxyz')
    print(key)
    print(nonce)
    print(tag)
    print(data)
    decryptor = Decryptor(key, nonce, tag, data)
    print(decryptor.decrypt())


if __name__ == '__main__':
    test()
